use std::sync::Arc;

use async_trait::async_trait;
use db::dblogic::db::establish_connection;
use sea_orm::sqlx::types::chrono;
use sea_orm::*;

use crate::base::Task;
use crate::entity::{self, ActiveModel};
use crate::persistence::{PersistenceError, TaskPersistence, TaskStatus};

pub struct SeaOrmTaskPersistence {
    db: DatabaseConnection,
    deserialize_fn: Box<dyn Fn(&str, &str) -> Result<Box<dyn Task>, String> + Send + Sync>,
}

impl SeaOrmTaskPersistence {
    pub async fn new(
        deserialize_fn: impl Fn(&str, &str) -> Result<Box<dyn Task>, String> + Send + Sync + 'static
    ) -> Self {
        Self { 
            db: establish_connection().await,
            deserialize_fn: Box::new(deserialize_fn),
        }
    }
}

#[async_trait]
impl TaskPersistence for SeaOrmTaskPersistence {
    async fn save_task(&self, task_id: &str, task: &Arc<Box<dyn Task>>, status: TaskStatus) -> Result<(), PersistenceError> {
        let task_data = task.serialize().map_err(|e| PersistenceError::SaveError(e.to_string()))?;
        
        let executor_status = ActiveModel {
            task_id: Set(task_id.to_string()),
            task_type: Set(task.type_name().to_string()),
            task_data: Set(task_data),
            status: Set(format!("{:?}", status)),
            created_at: Set(chrono::Utc::now().into()),
            updated_at: Set(chrono::Utc::now().into()),
        };

        entity::Entity::insert(executor_status)
            .exec(&self.db)
            .await
            .map_err(|e| PersistenceError::SaveError(e.to_string()))?;

        Ok(())
    }

    async fn remove_task(&self, task_id: &str) -> Result<(), PersistenceError> {
        entity::Entity::delete_by_id(task_id)
            .exec(&self.db)
            .await
            .map_err(|e| PersistenceError::RemoveError(e.to_string()))?;

        Ok(())
    }

    async fn update_task_status(&self, task_id: &str, status: TaskStatus) -> Result<(), PersistenceError> {
        let executor_status = entity::Entity::find_by_id(task_id)
            .one(&self.db)
            .await
            .map_err(|e| PersistenceError::UpdateError(e.to_string()))?
            .ok_or_else(|| PersistenceError::NotFound(task_id.to_string()))?;

        let mut executor_status: ActiveModel = executor_status.into();
        executor_status.status = Set(format!("{:?}", status));
        executor_status.updated_at = Set(chrono::Utc::now().into());

        entity::Entity::update(executor_status)
            .exec(&self.db)
            .await
            .map_err(|e| PersistenceError::UpdateError(e.to_string()))?;

        Ok(())
    }

    async fn get_all_tasks(&self) -> Result<Vec<(String, Arc<Box<dyn Task>>, TaskStatus)>, PersistenceError> {
        let executor_statuses = entity::Entity::find()
            .all(&self.db)
            .await
            .map_err(|e| PersistenceError::SaveError(e.to_string()))?;

        let mut tasks = Vec::new();
        for status in executor_statuses {
            let task_type = status.task_type.as_str();
            let task_data = status.task_data.as_str();
            let task_status = status.status.parse().map_err(|_| {
                PersistenceError::SaveError("Invalid task status".to_string())
            })?;

            let task = (self.deserialize_fn)(task_type, task_data)
                .map_err(|e| PersistenceError::SaveError(e))?;

            tasks.push((status.task_id, Arc::new(task), task_status));
        }

        Ok(tasks)
    }
}